package com.liru.chat_demo.controller;

import com.liru.chat_demo.model.FunctionCallRequest;
import com.liru.chat_demo.model.FunctionSchema;
import com.liru.chat_demo.service.Assistant;
import com.liru.chat_demo.service.McpClientService;
import jakarta.annotation.Resource;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@RestController
public class AiController {
    @Resource
    private Assistant assistant;
    private final McpClientService mcpClientService;
    private final Map<String, List<String>> sessionHistory = new ConcurrentHashMap<>(); // 存储会话历史

    public AiController(McpClientService mcpClientService) {
        this.mcpClientService = mcpClientService;
    }

    @RequestMapping(value = "/chat1", produces = "text/event-stream;charset=UTF-8")
    public Flux<String> chat1(@RequestParam String userMessage) {
        String id = "default";
        return assistant.chat(id, userMessage);
    }

    @GetMapping(value = "/chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chat(@RequestParam String userMessage, @RequestParam(required = false) String sessionId) {
        sessionHistory.computeIfAbsent(sessionId, k -> new ArrayList<>()).add("You: " + userMessage);
        sessionHistory.forEach((k, v) -> System.out.println(k + ": " + v));
        if (sessionId == null || sessionId.isBlank())
            sessionId = "default";
        return assistant.chat(sessionId, userMessage);
    }
    

    /**
     * 增强版聊天端点，集成MCP工具
     * 该方法提供了一个可以使用MCP工具的聊天接口
     *
     * @param userMessage 用户消息
     * @param sessionId   会话ID（可选）
     * @return 作为流返回的聊天响应
     */
    @GetMapping(value = "/chat2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public String chat2(@RequestParam String userMessage, @RequestParam(required = false) String sessionId) {
        // 将用户消息存储到会话历史中
        if (sessionId == null || sessionId.isBlank()) {
            sessionId = "default";
        }
        sessionHistory.computeIfAbsent(sessionId, k -> new ArrayList<>()).add("You: " + userMessage);
        
        // 获取LLM可用的函数
        List<FunctionSchema> functions = mcpClientService.getFunctionSchemas();
        
        // 准备包含函数的上下文
        Map<String, Object> context = new HashMap<>();
        context.put("functions", functions);
        context.put("userMessage", userMessage);
        
        // 使用LLM处理（为演示而简化）
        // 在实际实现中，您需要:
        // 1. 将用户消息和函数发送给LLM
        // 2. 获取LLM响应，可能包含函数调用
        // 3. 执行函数调用并将结果返回给LLM
        // 4. 获取最终的LLM响应
        
        // 在此演示中，我们将通过检查消息是否包含与可用工具相关的关键词来模拟这一过程，并直接调用它们
        
        // 检查消息是否包含与天气相关的关键词
        if (userMessage.toLowerCase().contains("weather") && 
            functions.stream().anyMatch(f -> f.getName().contains("Weather") || f.getName().contains("weather"))) {
            
            // 出于演示目的，如果可用则调用天气工具
            String andCallWeatherTool = findAndCallWeatherTool(functions, userMessage);
            return andCallWeatherTool;
        }
        
        // 如果没有匹配的工具，则回退到常规聊天
        return assistant.chat1(sessionId, userMessage);
    }
    
    /**
     * 查找并调用与天气相关的工具
     *
     * @param functions   可用的函数
     * @param userMessage 用户消息
     * @return 工具结果
     */
    private String findAndCallWeatherTool(List<FunctionSchema> functions, String userMessage) {
        // 查找与天气相关的工具
        FunctionSchema weatherTool = functions.stream()
                .filter(f -> f.getName().contains("Weather") || f.getName().contains("weather"))
                .findFirst()
                .orElse(null);
        
        if (weatherTool == null) {
            return "我找不到能帮助您处理请求的天气工具。";
        }
        
        // 为了演示目的，我们使用一些默认参数
        Map<String, Object> args = new HashMap<>();
        
        // 检查是否为getWeatherForecastByLocation工具
        if (weatherTool.getName().equals("getWeatherForecastByLocation")) {
            args.put("latitude", "47.6062");
            args.put("longitude", "-122.3321");
        } 
        // 检查是否为getAlerts工具
        else if (weatherTool.getName().equals("getAlerts")) {
            args.put("state", "NY");
        }
        
        // 调用工具
        FunctionCallRequest request = new FunctionCallRequest(weatherTool.getName(), args);
        return mcpClientService.callTool(request.getToolName(), request.getArguments());
    }
}
